"""excel merger
    examples: python main.py old.xlsx new.xlsx
    references:
    1. https://segmentfault.com/a/1190000018177573
    2. https://blog.csdn.net/qq_34377830/article/details/81872568

"""

from cmath import nan
from email.header import Header
from math import comb
from numpy import equal
import pandas as pd
from pandas import DataFrame,ExcelWriter
import os

date = '2022.01.23'

def getTotalData(topdir):
    total_fname = ''
    for fname in os.listdir(topdir):
        if fname.endswith('.xlsx'):
            total_fname = fname
            break
    old_fiveclass_data = pd.read_excel(os.path.join(topdir, fname), 0)
    old_hesuan_data = pd.read_excel(os.path.join(topdir, fname), '核酸')
    # total_1: 前一天5类数据，[乡镇，摸排累计，返乡累计]
    total_1 = old_fiveclass_data.iloc[3:len(old_fiveclass_data)-3,[0,18,20]].values.tolist()
    # total_2: 前一天核酸数据，[累计返武人员数]
    total_2 = old_hesuan_data.iloc[3:len(old_hesuan_data)-1, [2]].values.tolist()

    total_data = []
    for start, end in zip(total_1, total_2):
        total_data.append(start + end)
    return total_data

def formatData(data):
    if isinstance(data, int):
        return data
    elif isinstance(data, str):
        return 0
    return 0

def getFiveClassData(topdir):
    # datas = [ ['乡镇', '摸排数-非高中风险', '摸排数-高中风险', '摸排数-总数', '返乡数-非高中风险', '返乡数-高中风险', '返乡数-总数' ] ]
    datas = []
    for fname in os.listdir(topdir):
        data = [ fname.split('.')[0] ]
        orig = pd.read_excel(os.path.join(topdir, fname), 0)
        mopai_low = orig.iloc[6, 3]
        mopai_high = orig.iloc[7, 3]
        mopai_sum = formatData(mopai_low) + formatData(mopai_high)
        data.extend([mopai_low, mopai_high, mopai_sum])

        fan_low = orig.iloc[6, 4]
        fan_high = orig.iloc[7, 4]
        fan_sum = formatData(fan_low) + formatData(fan_high)

        data.extend([fan_low, fan_high, fan_sum])
        datas.append(data)
    return datas

def createTableHeader():
    datas = [
        ["", "前一天", "", "", "今日", "", "", "", "", "", "第5类", "", "", ""],
        ["","摸排","返乡","累计返武人员数","摸排数", "","","返乡数","", "","摸排","","返乡",""],
        ["乡镇","累计","累计","累计","非高中风险","高中风险","总数","非高中风险","高中风险","总数","当日新增","累计","当日新增","累计"]]
    return datas


def combineDatas(old, new):
    old_seqs = [ i[0] for i in old]
    new_seqs = [ i[0] for i in new]
    
    combined_datas = []
    for i,name in enumerate(old_seqs):
        ind = [j for j,x in enumerate(new_seqs) if x.find(name)!=-1]
        data = old[i]
        if not ind:
            data.extend([nan] * 10)
        else:
            data.extend(new[ind[0]][1:])
            # 摸排数当日新增，累计，返乡数当日新增，累计
            mopai_total = data[6]
            mopai_increase = mopai_total - data[1]
            fan_total = data[9]
            fan_increase = fan_total - data[2]
            data.extend([mopai_increase, mopai_total, fan_increase, fan_total])
        combined_datas.append(data)
    return combined_datas


def main(topdir):
    tbl_header = createTableHeader()
    old_datas = getTotalData(topdir)
    new_datas = getFiveClassData(topdir+'/5类/')
    combined_datas = combineDatas(old_datas, new_datas)

    total_datas = DataFrame(tbl_header + combined_datas)
    print(total_datas)

    writer = ExcelWriter(date + ' 5类汇总.xlsx')
    total_datas.to_excel(writer, '5类汇总', index=None, header=None)
    writer.save()

main(date)
